package device

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"time"

	"golang.org/x/crypto/chacha20poly1305"
)

func (peer *Peer) sendHandshakeInitiation(isRetry bool) error {
	if !isRetry {
		peer.timers.handshakeAttempts.Store(0)
	}
	peer.handshake.RLock()
	if time.Since(peer.handshake.lastSentHandshake) < time.Second*5 {
		peer.handshake.RUnlock()
		return nil
	}
	peer.handshake.RUnlock()
	peer.handshake.Lock()
	if time.Since(peer.handshake.lastSentHandshake) < time.Second*5 {
		peer.handshake.Unlock()
		return nil
	}
	peer.handshake.lastSentHandshake = time.Now()
	peer.handshake.Unlock()

	fmt.Println(" - Sending handshake initiation")
	msg, err := peer.device.createMessageInitiation(peer)
	if err != nil {
		fmt.Printf("Failed to create initiation message: %v\n", err)
		return err
	}

	var buf [148]byte
	writer := bytes.NewBuffer(buf[:0])
	binary.Write(writer, binary.LittleEndian, msg)
	packet := writer.Bytes()
	peer.cookieGenerator.addMacs(packet)
	peer.device.udpConn.WriteToUDP(packet, peer.endpoint.addr)
	peer.timersAnyAuthenticatedPacketTraversal()
	peer.timersHandshakeInitiated()
	return nil
}
func (peer *Peer) sendHandshakeResponse() error {
	peer.handshake.Lock()
	peer.handshake.lastSentHandshake = time.Now()
	peer.handshake.Unlock()
	response, err := peer.device.createMessageResponse(peer)
	if err != nil {
		fmt.Printf("Failed to create response message: %v\n", err)
		return err
	}

	var buf [92]byte
	writer := bytes.NewBuffer(buf[:0])
	binary.Write(writer, binary.LittleEndian, response)
	packet := writer.Bytes()
	peer.cookieGenerator.addMacs(packet)
	err = peer.beginSymmetricSession()
	if err != nil {
		fmt.Println("Failed to derive keypair for send handshake response")
		return err
	}
	peer.device.udpConn.WriteToUDP(packet, peer.endpoint.addr)
	return nil
}

func (peer *Peer) sendKeepalive() {
	buff := make([]byte, 32)
	header := buff[:16]
	fieldType := header[0:4]
	fieldReceiver := header[4:8]
	fieldNonce := header[8:16]
	keyPair := peer.keypairs.Current()
	if keyPair == nil {
		fmt.Println("没有可用的密钥对")
		return
	}

	counter := keyPair.sendNonce.Add(1) - 1
	binary.LittleEndian.PutUint32(fieldType, 4)
	binary.LittleEndian.PutUint32(fieldReceiver, keyPair.remoteIndex)
	binary.LittleEndian.PutUint64(fieldNonce, counter)
	var nonce [chacha20poly1305.NonceSize]byte
	binary.LittleEndian.PutUint64(nonce[4:], counter)
	emptyArea := make([]byte, 0)
	buff = keyPair.send.Seal(
		header,
		nonce[:],
		emptyArea,
		nil,
	)
	peer.endpoint.RLock()
	peer.device.udpConn.WriteToUDP(buff, peer.endpoint.addr)
	peer.endpoint.RUnlock()
	peer.timersAnyAuthenticatedPacketTraversal()
}

func (device *Device) tansportDataSendHandle() {
	device.Wg.Add(1)
	defer device.Wg.Done()
	var paddingZeros [16]byte
	var nonce [chacha20poly1305.NonceSize]byte
	for {
		select {
		case <-device.Done:
			return
		case packet_ := <-device.queue.send:
			buffer := *device.pool.Get()
			{
				peer := device.AllowedIPs.find(packet_.endpoint.dstIP)
				if peer == nil {
					fmt.Println("Received packet, but no peer found", packet_.endpoint.dstIP)
					goto skip
				}
				keyPair := peer.keypairs.Current()
				if keyPair == nil {
					fmt.Println("没有可用的密钥对")
				}
				if keyPair == nil || keyPair.sendNonce.Load() >= RejectAfterMessages || time.Since(keyPair.created) >= RejectAfterTime {
					peer.sendHandshakeInitiation(false)
					goto skip
				}
				counter := keyPair.sendNonce.Add(1) - 1
				if counter >= RejectAfterMessages {
					keyPair.sendNonce.Store(RejectAfterMessages)
					goto skip
				}
				header := buffer[:16]
				fieldType := header[0:4]
				fieldReceiver := header[4:8]
				fieldNonce := header[8:16]
				binary.LittleEndian.PutUint32(fieldType, 4)
				binary.LittleEndian.PutUint32(fieldReceiver, keyPair.remoteIndex)
				binary.LittleEndian.PutUint64(fieldNonce, counter)

				// pad content to multiple of 16
				paddingSize := calculatePaddingSize(len(packet_.packet), device.Tun.Mtu())
				packet_.packet = append(packet_.packet, paddingZeros[:paddingSize]...)

				// encrypt content and release to consumer
				binary.LittleEndian.PutUint64(nonce[4:], counter)
				packet_.packet = keyPair.send.Seal(
					header,
					nonce[:],
					packet_.packet,
					nil,
				)
				peer.endpoint.RLock()
				peer.device.udpConn.WriteToUDP(packet_.packet, peer.endpoint.addr)
				peer.endpoint.RUnlock()
				peer.timersAnyAuthenticatedPacketTraversal()
				peer.keepKeyFreshSending()
				peer.timersDataSent()
				goto skip
			}
		skip:
			// 回收到对象池
			device.pool.Put(packet_.buffer)
			device.pool.Put(&buffer)
		}
	}
}
func calculatePaddingSize(packetSize, mtu int) int {
	lastUnit := packetSize
	PaddingMultiple := 16
	if mtu == 0 {
		return ((lastUnit + PaddingMultiple - 1) & ^(PaddingMultiple - 1)) - lastUnit
	}
	if lastUnit > mtu {
		lastUnit %= mtu
	}
	paddedSize := ((lastUnit + PaddingMultiple - 1) & ^(PaddingMultiple - 1))
	if paddedSize > mtu {
		paddedSize = mtu
	}
	return paddedSize - lastUnit
}

func (peer *Peer) keepKeyFreshSending() {
	keypair := peer.keypairs.Current()
	if keypair == nil {
		return
	}
	nonce := keypair.sendNonce.Load()
	if nonce > RekeyAfterMessages || (keypair.isInitiator && time.Since(keypair.created) > RekeyAfterTime) {
		peer.sendHandshakeInitiation(false)
	}
}
